#include<cmath>
#include<algorithm>
#include"system_hf.h"
#include"help.h"
using std::abs;
using std::min;
using std::max;

void System_Tz_HF::setupOrbitals()
{
  hfsolver.init();
  hfsolver.iter();
  hfsolver.calEnergy();
  HFenergy=hfsolver.Energy;
  Orbitals=hfsolver.HFOrbitals;
  setupGroups();
  //set fermi surface
  FermiSurface=0;
  for(int i=0;i<Orbitals.size();i++)
    {
      if(Orbitals[i].state) break;
      FermiSurface=i;
    }
}

/*
  dim of MatEle=4,the 4 double are:
      0  interaction
      1  (p1 \dot p2)/(m hbar omega) + m omega r1 \dot r2 / hbar
      2  (r1-r2)^2 m omega/(2 hbar)
      3  (p1 \dot p2)/(m hbar omega)
*/
const int System_Tz_HF::MatEledim=System_TzType::MatEledim;
void System_Tz_HF::setupTwoBodyMat()
{
  setupIndiceIn2BStates();
  int totalChannels=TwoBodyChannels.size();
  TwoBodyMat.resize(totalChannels);
  #pragma omp parallel for
  for(int channel=0;channel<totalChannels;channel++)
    {
      int num2BInChannel=TwoBodyChannels[channel].size();
      TwoBodyMat[channel].clear();
      TwoBodyMat[channel].resize( (num2BInChannel*(num2BInChannel+1))/2 );
      for(int bra=0;bra<num2BInChannel;bra++)
	for(int ket=bra;ket<num2BInChannel;ket++)
	  {
	    int ab=TwoBodyChannels[channel][bra];
	    int cd=TwoBodyChannels[channel][ket];
	    int a=TwoBodyStates[ab].a;
	    int b=TwoBodyStates[ab].b;
	    int c=TwoBodyStates[cd].a;
	    int d=TwoBodyStates[cd].b;
	    int J=TwoBodyStates[ab].J;
	    TwoBodyMatEle MatEle(MatEledim);
	    get2BmatOnNeed(a,b,c,d,J,MatEle);
	    set2Bmat(channel,bra,ket,MatEle);
	  }
    }
}

void System_Tz_HF::get2BmatOnNeed(int a,int b,int c,int d,int J,TwoBodyMatEle & MatEle) const
{
  MatEle.resize(MatEledim);
  MatEle.setZero();
  int jja=Orbitals[a].jj;
  int jjb=Orbitals[b].jj;
  int jjc=Orbitals[c].jj;
  int jjd=Orbitals[d].jj;
  int Tzab=(Orbitals[a].tz+Orbitals[b].tz)/2;
  int Tzcd=(Orbitals[c].tz+Orbitals[d].tz)/2;

  if(a==b && J%2) return;
  if(c==d && J%2) return;
  if(Tzab!=Tzcd) return;//isospin projection check
  if((Orbitals[a].l+Orbitals[b].l+Orbitals[c].l+Orbitals[d].l)%2) return;//parity check
  int Jmin=max(abs(jja-jjb)/2,abs(jjc-jjd)/2);
  int Jmax=min((jja+jjb)/2,(jjc+jjd)/2);
  if(Jmin>Jmax) return;//angular momenta check

  for(int ia=0;ia<Orbitals[a].Coeff.size();ia++)
    for(int ib=0;ib<Orbitals[b].Coeff.size();ib++)
      for(int ic=0;ic<Orbitals[c].Coeff.size();ic++)
	for(int id=0;id<Orbitals[d].Coeff.size();id++)
	  {
	    System_TzType::TwoBodyMatEle Temp;
	    int iia=pSystem_Tz->Groups[Orbitals[a].GroupIndex][ia];
	    int iib=pSystem_Tz->Groups[Orbitals[b].GroupIndex][ib];
	    int iic=pSystem_Tz->Groups[Orbitals[c].GroupIndex][ic];
	    int iid=pSystem_Tz->Groups[Orbitals[d].GroupIndex][id];
	    pSystem_Tz->get2BmatAt(iia,iib,iic,iid,J,Temp);
	    if(Temp.empty()) continue;
	    if(iia==iib) Temp*=sqrt(2.);
	    if(iic==iid) Temp*=sqrt(2.);
	    MatEle+=Temp*(Orbitals[a].Coeff[ia]*Orbitals[b].Coeff[ib]*Orbitals[c].Coeff[ic]*Orbitals[d].Coeff[id]);
	  }
  if(a==b) MatEle*=sqrt(2.)/2;
  if(c==d) MatEle*=sqrt(2.)/2;
}

// void System_Tz_HF::setupTwoBodyMat()
// {
//    setupIndiceIn2BStates();
//   int totalChannels=TwoBodyChannels.size();
//   TwoBodyMat.resize(totalChannels);
//   for(int channel=0;channel<totalChannels;channel++)
//     {
//       int num2BInChannel=TwoBodyChannels[channel].size();
//TwoBodyMat[channel].clear();
//       TwoBodyMat[channel].resize( (num2BInChannel*(num2BInChannel+1))/2 );
      
//       int Vdim=pSystem_Tz->TwoBodyChannels[channel].size();
//       Mat2B Vmat(Vdim,Vdim);
//       //setup Vmat
//       for(int i=0;i<Vdim;i++)
// 	{
// 	  for(int j=i;j<Vdim;j++)
// 	    {
// 	      pSystem_Tz->get2Bmat(channel,i,j,Vmat(i,j));
// 	      Vmat(j,i)=Vmat(i,j);
// 	    }
// 	}
//       //setup TransMat
//       Mat TransMat(Vdim,num2BInChannel);
//       for(int i=0;i<Vdim;i++)
// 	for(int j=0;j<num2BInChannel;j++)
// 	  {
// 	    int ab=pSystem_Tz->TwoBodyChannels[channel][i];
// 	    int alphabeta=TwoBodyChannels[channel][j];
// 	    int a=pSystem_Tz->TwoBodyStates[ab].a;
// 	    int b=pSystem_Tz->TwoBodyStates[ab].b;
// 	    int alpha=TwoBodyStates[alphabeta].a;
// 	    int beta=TwoBodyStates[alphabeta].b;
// 	    //debug: TwoBodyState ab and alphabeta should have same quantum number
// 	    DataType coef=overlap(a,alpha)*overlap(b,beta);
// 	    int jjalpha=Orbitals[alpha].jj;
// 	    int jjbeta=Orbitals[beta].jj;
// 	    int J=TwoBodyStates[alphabeta].J;
// 	    coef-=phase( (jjalpha+jjbeta)/2 -J )*overlap(a,beta)*overlap(b,alpha);
// 	    if(a==b) coef/=sqrt(2.);
// 	    if(alpha==beta) coef/=sqrt(2.);
// 	    TransMat(i,j)=coef;
// 	  }
//       cout<<channel<<endl;
//       //transform  and write into TwoBodyMat
//       for(int i=0;i<num2BInChannel;i++)
//       	for(int j=i;j<num2BInChannel;j++)
//       	  {
//       	    TwoBodyMatEle VHF(MatEledim);
//       	    for(int k=0;k<Vdim;k++)
//       	      for(int l=0;l<Vdim;l++)
//       		{
// 		  if(TransMat(l,j)==0||TransMat(k,i)==0) continue;
// 		  if(Vmat(k,l).empty()) continue;
//       		  VHF+=Vmat(k,l)*(TransMat(l,j)*TransMat(k,i));
//       		}
//       	    set2Bmat(channel,i,j,VHF);
//       	  }
//     }
// }
System_Tz_HF::DataType System_Tz_HF::get1B(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	    val+=pSystem_Tz->get1B(realA,iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
}

System_Tz_HF::DataType System_Tz_HF::get1B_p2(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];	    
	    val+=pSystem_Tz->get1B_p2(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
}
System_Tz_HF::DataType System_Tz_HF::get1B_p2_plus_r2(int bra,int ket) const
{
  DataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	  {
	    int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	    int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	    if(iibra!=iiket) continue;
	    val+=pSystem_Tz->Orbitals[iibra].e*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	    //val+=pSystem_Tz->get1B_p2_plus_r2(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	  }
    }
  return val;
}


System_Tz_HF::DataType System_Tz_HF::rL(int bra,int ket,int L) const
{
  DataType val=0;
  for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
    for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
      {
	int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
	int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
	val+=pSystem_Tz->rL(iibra,iiket,L)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
      }
  return val;
}
System_Tz_HF::DataType System_Tz_HF::Q(int p,int h,int L,int isospin) const
{
  int Tzh=Orbitals[h].tz;
  int Tzp=Orbitals[p].tz;
  if(Tzh!=Tzp) return 0;  
  int jjp=Orbitals[p].jj;
  int jjh=Orbitals[h].jj;
  int lp=Orbitals[p].l;
  int lh=Orbitals[h].l;
  DataType temp,fac;
  if(L==0)
    {
      temp=rL(p,h,2);
      temp*=reducedYMatEle(lp,jjp,lh,jjh,0)* sqrt(4*Pi);
    }
  else
    {
      temp=rL(p,h,L);
      temp*=reducedYMatEle(lp,jjp,lh,jjh,L);
    }
  fac=isospin?Tzh:1.0;
  temp*=fac;
  return temp;
}

System_Tz_HF::DataType System_Tz_HF::Q1_v(int p,int h) const
{
  int Tzh=Orbitals[h].tz;
  int Tzp=Orbitals[p].tz;
  if(Tzh!=Tzp) return 0;
  DataType temp=rL(p,h,1);
  int jjp=Orbitals[p].jj;
  int jjh=Orbitals[h].jj;
  int lp=Orbitals[p].l;
  int lh=Orbitals[h].l;
  temp*=reducedYMatEle(lp,jjp,lh,jjh,1);
  temp*=(Tzh*0.5-0.5+Z*1.0/realA);
  return temp;
}
System_Tz_HF::DataType System_Tz_HF::Q1_s(int p,int h,DataType Rms) const
{
  int Tzh=Orbitals[h].tz;
  int Tzp=Orbitals[p].tz;
  if(Tzh!=Tzp) return 0;
  DataType temp=rL(p,h,3)-rL(p,h,1)*Rms*5./3;

  int jjp=Orbitals[p].jj;
  int jjh=Orbitals[h].jj;
  int lp=Orbitals[p].l;
  int lh=Orbitals[h].l;
  temp*=reducedYMatEle(lp,jjp,lh,jjh,1);
  return temp;
}
